import faiss
import numpy as np
import bin_io


def ip_gnd(base, query, k):
    base_dim = base.shape[1]
    index = faiss.IndexFlatIP(base_dim)
    index.add(base)
    gnd_distance, gnd_idx = index.search(query, k)
    return gnd_idx, gnd_distance


if __name__ == '__main__':
    base, n_base, dim = bin_io.fbin_read('/home/bianzheng/NIPS-Competition/data/text-to-image-40M/base.fbin')
    print('read base')
    query, n_query, dim = bin_io.fbin_read('/home/bianzheng/NIPS-Competition/data/text-to-image/query.public.100K.fbin')
    query = query[:10000, :]
    bin_io.fbin_write('/home/bianzheng/NIPS-Competition/data/text-to-image-40M/query.fbin', query)

    idx_, distance_ = ip_gnd(base, query, 10)
    bin_io.ibin_write('/home/bianzheng/NIPS-Competition/data/text-to-image-40M/gnd-top10.ibin', idx_)
